import os
import os.path as osp
import ctypes
import numpy as np
import OpenGL.GL as gl
import glm
import pyassimp
from PIL import Image

MAX_BONE_INFLUENCE = 4

class Vertex:
    def __init__(self, position, normal, texCoords, tangent, bitangent, boneIDs, weights):
        self.position = position
        self.normal = normal
        self.texCoords = texCoords
        self.tangent = tangent
        self.bitangent = bitangent
        self.boneIDs = boneIDs
        self.weights = weights

class Texture:
    def __init__(self, id, type, path):
        self.id = id
        self.type = type
        self.path = path

class Mesh:
    def __init__(self, vertices, indices, textures):
        self.vertices = vertices
        self.indices = indices
        self.textures = textures
        self.VAO = gl.glGenVertexArrays(1)
        self.setup_mesh()

    def setup_mesh(self):
        self.VBO = gl.glGenBuffers(1)
        self.EBO = gl.glGenBuffers(1)
        
        gl.glBindVertexArray(self.VAO)

        gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.VBO)
        vertex_data = np.array([vertex_data for vertex in self.vertices for vertex_data in (
            vertex.position.tolist() + vertex.normal.tolist() + vertex.texCoords.tolist() +
            vertex.tangent.tolist() + vertex.bitangent.tolist() + vertex.boneIDs + vertex.weights)], dtype=np.float32)
        gl.glBufferData(gl.GL_ARRAY_BUFFER, vertex_data.nbytes, vertex_data, gl.GL_STATIC_DRAW)

        gl.glBindBuffer(gl.GL_ELEMENT_ARRAY_BUFFER, self.EBO)
        index_data = np.array(self.indices, dtype=np.uint32)
        gl.glBufferData(gl.GL_ELEMENT_ARRAY_BUFFER, index_data.nbytes, index_data, gl.GL_STATIC_DRAW)

        stride = vertex_data.itemsize * 18  # 3 + 3 + 2 + 3 + 3 + 4 + 4

        # vertex positions
        gl.glEnableVertexAttribArray(0)
        gl.glVertexAttribPointer(0, 3, gl.GL_FLOAT, gl.GL_FALSE, stride, ctypes.c_void_p(0))
        # vertex normals
        gl.glEnableVertexAttribArray(1)
        gl.glVertexAttribPointer(1, 3, gl.GL_FLOAT, gl.GL_FALSE, stride, ctypes.c_void_p(12))
        # vertex texture coords
        gl.glEnableVertexAttribArray(2)
        gl.glVertexAttribPointer(2, 2, gl.GL_FLOAT, gl.GL_FALSE, stride, ctypes.c_void_p(24))
        # vertex tangent
        gl.glEnableVertexAttribArray(3)
        gl.glVertexAttribPointer(3, 3, gl.GL_FLOAT, gl.GL_FALSE, stride, ctypes.c_void_p(32))
        # vertex bitangent
        gl.glEnableVertexAttribArray(4)
        gl.glVertexAttribPointer(4, 3, gl.GL_FLOAT, gl.GL_FALSE, stride, ctypes.c_void_p(44))
        # bone ids
        gl.glEnableVertexAttribArray(5)
        gl.glVertexAttribIPointer(5, 4, gl.GL_INT, stride, ctypes.c_void_p(56))
        # weights
        gl.glEnableVertexAttribArray(6)
        gl.glVertexAttribPointer(6, 4, gl.GL_FLOAT, gl.GL_FALSE, stride, ctypes.c_void_p(72))

        gl.glBindVertexArray(0)

    def draw(self, shader):
        diffuseNr = 1
        specularNr = 1
        normalNr = 1
        heightNr = 1

        for i, texture in enumerate(self.textures):
            gl.glActiveTexture(gl.GL_TEXTURE0 + i)
            name = texture.type
            number = ''

            if name == 'texture_diffuse':
                number = str(diffuseNr)
                diffuseNr += 1
            elif name == 'texture_specular':
                number = str(specularNr)
                specularNr += 1
            elif name == 'texture_normal':
                number = str(normalNr)
                normalNr += 1
            elif name == 'texture_height':
                number = str(heightNr)
                heightNr += 1

            gl.glUniform1i(gl.glGetUniformLocation(shader.program, (name + number).encode('utf-8')), i)
            gl.glBindTexture(gl.GL_TEXTURE_2D, texture.id)

        gl.glBindVertexArray(self.VAO)
        gl.glDrawElements(gl.GL_TRIANGLES, len(self.indices), gl.GL_UNSIGNED_INT, None)
        gl.glBindVertexArray(0)
        gl.glActiveTexture(gl.GL_TEXTURE0)

class Model:
    def __init__(self, path, gamma=False):
        self.textures_loaded = []
        self.meshes = []
        self.directory = ''
        self.gammaCorrection = gamma
        self.load_model(path)

    def draw(self, shader):
        for mesh in self.meshes:
            mesh.draw(shader)

    def load_model(self, path):
        path = osp.abspath(path)
        assert osp.exists(path)
        # scene = pyassimp.load(path)
        scene = pyassimp.load(path, processing=pyassimp.postprocess.aiProcess_Triangulate |
                                   pyassimp.postprocess.aiProcess_GenSmoothNormals |
                                   pyassimp.postprocess.aiProcess_FlipUVs |
                                   pyassimp.postprocess.aiProcess_CalcTangentSpace)
        if not scene or not scene.meshes:
            print("ERROR::ASSIMP:: ", pyassimp.get_error())
            return

        self.directory = path.rsplit('/', 1)[0]
        self.process_node(scene.rootnode, scene)
        pyassimp.release(scene)

    def process_node(self, node, scene):
        for mesh in node.meshes:
            self.meshes.append(self.process_mesh(mesh, scene))
        
        for child in node.children:
            self.process_node(child, scene)

    def process_mesh(self, mesh, scene):
        vertices = []
        indices = []
        textures = []

        for i in range(mesh.numvertices):
            position = glm.vec3(mesh.vertices[i][0], mesh.vertices[i][1], mesh.vertices[i][2])
            normal = glm.vec3(mesh.normals[i][0], mesh.normals[i][1], mesh.normals[i][2]) if mesh.has_normals else glm.vec3(0.0, 0.0, 0.0)
            texCoords = glm.vec2(mesh.texturecoords[0][i][0], mesh.texturecoords[0][i][1]) if mesh.texturecoords else glm.vec2(0.0, 0.0)
            tangent = glm.vec3(mesh.tangents[i][0], mesh.tangents[i][1], mesh.tangents[i][2]) if mesh.has_tangents_and_bitangents else glm.vec3(0.0, 0.0, 0.0)
            bitangent = glm.vec3(mesh.bitangents[i][0], mesh.bitangents[i][1], mesh.bitangents[i][2]) if mesh.has_tangents_and_bitangents else glm.vec3(0.0, 0.0, 0.0)
            boneIDs = [0] * MAX_BONE_INFLUENCE
            weights = [0.0] * MAX_BONE_INFLUENCE

            vertices.append(Vertex(position, normal, texCoords, tangent, bitangent, boneIDs, weights))

        for face in mesh.faces:
            indices.extend(face.indices)
        
        material = scene.materials[mesh.materialindex]
        textures.extend(self.load_material_textures(material, pyassimp.assimp_aiTextureType_DIFFUSE, 'texture_diffuse'))
        textures.extend(self.load_material_textures(material, pyassimp.assimp_aiTextureType_SPECULAR, 'texture_specular'))
        textures.extend(self.load_material_textures(material, pyassimp.assimp_aiTextureType_HEIGHT, 'texture_normal'))
        textures.extend(self.load_material_textures(material, pyassimp.assimp_aiTextureType_AMBIENT, 'texture_height'))

        return Mesh(vertices, indices, textures)

    def load_material_textures(self, mat, type, typeName):
        textures = []
        for i in range(mat.GetTextureCount(type)):
            str = mat.GetTexture(type, i)
            skip = False
            for j in range(len(self.textures_loaded)):
                if self.textures_loaded[j].path == str:
                    textures.append(self.textures_loaded[j])
                    skip = True
                    break
            if not skip:
                texture = Texture(TextureFromFile(str, self.directory), typeName, str)
                textures.append(texture)
                self.textures_loaded.append(texture)
        return textures

def TextureFromFile(path, directory, gamma=False):
    filename = f"{directory}/{path}"
    textureID = gl.glGenTextures(1)
    image = Image.open(filename)
    image = image.transpose(Image.FLIP_TOP_BOTTOM)
    img_data = np.array(list(image.getdata()), np.uint8)

    gl.glBindTexture(gl.GL_TEXTURE_2D, textureID)
    if image.mode == "RGB":
        gl.glTexImage2D(gl.GL_TEXTURE_2D, 0, gl.GL_RGB, image.width, image.height, 0, gl.GL_RGB, gl.GL_UNSIGNED_BYTE, img_data)
    elif image.mode == "RGBA":
        gl.glTexImage2D(gl.GL_TEXTURE_2D, 0, gl.GL_RGBA, image.width, image.height, 0, gl.GL_RGBA, gl.GL_UNSIGNED_BYTE, img_data)

    gl.glGenerateMipmap(gl.GL_TEXTURE_2D)
    gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_S, gl.GL_REPEAT)
    gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_T, gl.GL_REPEAT)
    gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR_MIPMAP_LINEAR)
    gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR)

    return textureID
